import torch
import torch.nn as nn
import torch.nn.functional as functional
import math


class Bottleneck(nn.Module):
    def __init__(self, nChannels, growthRate):
        super(Bottleneck, self).__init__()
        interChannels = 4 * growthRate
        self.bn1 = nn.BatchNorm1d(nChannels)
        self.conv1 = nn.Conv1d(nChannels, interChannels, kernel_size=1, bias=False)
        self.bn2 = nn.BatchNorm1d(interChannels)
        self.conv2 = nn.Conv1d(interChannels, growthRate, kernel_size=3, padding=1, bias=False)

    def forward(self, x):
        out = self.conv1(functional.relu(self.bn1(x)))
        out = self.conv2(functional.relu(self.bn2(out)))
        out = torch.cat((x, out), 1)
        return out


class SingleLayer(nn.Module):
    def __init__(self, nChannels, growthRate):
        super(SingleLayer, self).__init__()
        self.bn1 = nn.BatchNorm1d(nChannels)
        self.conv1 = nn.Conv1d(nChannels, growthRate, kernel_size=3, padding=1, bias=False)

    def forward(self, x):
        out = self.conv1(functional.relu(self.bn1(x)))
        out = torch.cat((x, out), 1)
        return out


class Transition(nn.Module):
    def __init__(self, nChannels, nOutChannels, down=False):
        super(Transition, self).__init__()
        self.bn1 = nn.BatchNorm1d(nChannels)
        self.conv1 = nn.Conv1d(nChannels, nOutChannels, kernel_size=1, bias=False)
        self.down = down

    def forward(self, x):
        out = self.conv1(functional.relu(self.bn1(x)))
        if self.down:
            out = functional.avg_pool1d(out, 2)
        return out


class ResidualUBlock(nn.Module):
    def __init__(self, out_ch, mid_ch, layers, downsampling=True):
        super(ResidualUBlock, self).__init__()
        self.downsample = downsampling
        K = 9  # Kernel size
        P = (K - 1) // 2  # Padding calculation

        # Initial convolutional layer
        self.conv1 = nn.Conv1d(in_channels=out_ch, out_channels=out_ch, kernel_size=K, padding=P, bias=False)
        self.bn1 = nn.BatchNorm1d(out_ch)

        self.encoders = nn.ModuleList()
        self.decoders = nn.ModuleList()

        for idx in range(layers):
            if idx == 0:
                self.encoders.append(nn.Sequential(
                    nn.Conv1d(in_channels=out_ch, out_channels=mid_ch, kernel_size=K, stride=2, padding=P, bias=False),
                    nn.BatchNorm1d(mid_ch),
                    nn.LeakyReLU()
                ))
            else:
                self.encoders.append(nn.Sequential(
                    nn.Conv1d(in_channels=mid_ch, out_channels=mid_ch, kernel_size=K, stride=2, padding=P, bias=False),
                    nn.BatchNorm1d(mid_ch),
                    nn.LeakyReLU()
                ))

            if idx == layers - 1:
                self.decoders.append(nn.Sequential(
                    nn.ConvTranspose1d(in_channels=mid_ch * 2, out_channels=out_ch, kernel_size=K, stride=2, padding=P, output_padding=1, bias=False),
                    nn.BatchNorm1d(out_ch),
                    nn.LeakyReLU()
                ))
            else:
                self.decoders.append(nn.Sequential(
                    nn.ConvTranspose1d(in_channels=mid_ch * 2, out_channels=mid_ch, kernel_size=K, stride=2, padding=P, output_padding=1, bias=False),
                    nn.BatchNorm1d(mid_ch),
                    nn.LeakyReLU()
                ))

        self.bottleneck = nn.Sequential(
            nn.Conv1d(in_channels=mid_ch, out_channels=mid_ch, kernel_size=K, padding=P, bias=False),
            nn.BatchNorm1d(mid_ch),
            nn.LeakyReLU()
        )

        if self.downsample:
            self.idfunc_0 = nn.AvgPool1d(kernel_size=2, stride=2)
            self.idfunc_1 = nn.Conv1d(in_channels=out_ch, out_channels=out_ch, kernel_size=1, bias=False)

    def forward(self, x):
        x_in = functional.leaky_relu(self.bn1(self.conv1(x)))
        out = x_in
        encoder_out = []
        for idx, layer in enumerate(self.encoders):
            # For 1D inputs, pad should be a 2-tuple (pad_left, pad_right)
            if out.size(-1) % 4 != 0:
                out = functional.pad(out, (1, 0))
            out = layer(out)
            encoder_out.append(out)

        out = self.bottleneck(out)

        for idx, layer in enumerate(self.decoders):
            out = layer(torch.cat([out, encoder_out[-1 - idx]], dim=1))

        out = out[..., :x_in.size(-1)]
        out += x_in

        if self.downsample:
            out = self.idfunc_0(out)
            out = self.idfunc_1(out)

        return out


def _make_dense(nChannels, growthRate, nDenseBlocks, bottleneck):
    layers = []
    for i in range(int(nDenseBlocks)):
        if bottleneck:
            layers.append(Bottleneck(nChannels, growthRate))
        else:
            layers.append(SingleLayer(nChannels, growthRate))
        nChannels += growthRate
    return nn.Sequential(*layers)


class ResNet1d(nn.Module):
    def __init__(self, nOUT, in_ch=12, out_ch=256, mid_ch=64):
        super(ResNet1d, self).__init__()
        self.conv = nn.Conv1d(in_channels=in_ch, out_channels=out_ch, kernel_size=15, padding=7, stride=2, bias=False)
        self.bn = nn.BatchNorm1d(out_ch)

        self.rub_0 = ResidualUBlock(out_ch=out_ch, mid_ch=mid_ch, layers=6)
        self.rub_1 = ResidualUBlock(out_ch=out_ch, mid_ch=mid_ch, layers=5)
        self.rub_2 = ResidualUBlock(out_ch=out_ch, mid_ch=mid_ch, layers=4)
        self.rub_3 = ResidualUBlock(out_ch=out_ch, mid_ch=mid_ch, layers=3)

        growthRate = 12
        reduction = 0.5
        nChannels = out_ch
        nDenseBlocks = 16

        self.dense1 = _make_dense(nChannels, growthRate=12, nDenseBlocks=nDenseBlocks, bottleneck=True)
        nChannels += nDenseBlocks * growthRate
        nOutChannels = int(math.floor(nChannels * reduction))
        self.trans1 = Transition(nChannels, nOutChannels)

        nChannels = nOutChannels
        self.dense2 = _make_dense(nChannels, growthRate=12, nDenseBlocks=nDenseBlocks, bottleneck=True)
        nChannels += nDenseBlocks * growthRate
        self.trans2 = Transition(nChannels, out_ch)

        self.mha = nn.MultiheadAttention(out_ch, 8)
        self.pool = nn.AdaptiveMaxPool1d(output_size=1)
        self.fc_1 = nn.Linear(out_ch, nOUT)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = functional.leaky_relu(self.bn(self.conv(x)))

        x = self.rub_0(x)
        x = self.rub_1(x)
        x = self.rub_2(x)
        x = self.rub_3(x)

        x = self.trans1(self.dense1(x))
        x = self.trans2(self.dense2(x))

        x = functional.dropout(x, p=0.5, training=self.training)

        x = x.permute(2, 0, 1)
        x, _ = self.mha(x, x, x)
        x = x.permute(1, 2, 0)

        x = self.pool(x).squeeze(2)
        x = self.fc_1(x)
        out = self.sigmoid(x)

        return out

class MetaResNet1d(nn.Module):
    def __init__(self, nOUT, in_ch=12, out_ch=256, mid_ch=64):
        super(MetaResNet1d, self).__init__()
        self.conv = nn.Conv1d(in_channels=in_ch, out_channels=out_ch, kernel_size=15, padding=7, stride=2, bias=False)
        self.bn = nn.BatchNorm1d(out_ch)

        self.rub_0 = ResidualUBlock(out_ch=out_ch, mid_ch=mid_ch, layers=8)
        self.rub_1 = ResidualUBlock(out_ch=out_ch, mid_ch=mid_ch, layers=7)
        self.rub_2 = ResidualUBlock(out_ch=out_ch, mid_ch=mid_ch, layers=6)
        self.rub_3 = ResidualUBlock(out_ch=out_ch, mid_ch=mid_ch, layers=5)

        growthRate = 12
        reduction = 0.5
        nChannels = out_ch
        nDenseBlocks = 16

        self.dense1 = _make_dense(nChannels, growthRate=12, nDenseBlocks=nDenseBlocks, bottleneck=True)
        nChannels += nDenseBlocks * growthRate
        nOutChannels = int(math.floor(nChannels * reduction))
        self.trans1 = Transition(nChannels, nOutChannels)

        nChannels = nOutChannels
        self.dense2 = _make_dense(nChannels, growthRate=12, nDenseBlocks=nDenseBlocks, bottleneck=True)
        nChannels += nDenseBlocks * growthRate
        self.trans2 = Transition(nChannels, out_ch)

        self.mha = nn.MultiheadAttention(out_ch, 8)
        self.pool = nn.AdaptiveMaxPool1d(output_size=1)

        # Meta features (gender, age) -> projection to match out_ch
        self.meta_fc = nn.Sequential(
            nn.Linear(7, out_ch),
            nn.ReLU(),
            nn.Dropout(0.3)
        )

        # Final fusion and prediction
        self.fc_1 = nn.Linear(out_ch * 2, nOUT)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, meta):
        """
        x: (B, C, T)
        meta: (B, 7)  # [gender, age]
        """
        x = functional.leaky_relu(self.bn(self.conv(x)))

        x = self.rub_0(x)
        x = self.rub_1(x)
        x = self.rub_2(x)
        x = self.rub_3(x)

        x = self.trans1(self.dense1(x))
        x = self.trans2(self.dense2(x))

        x = functional.dropout(x, p=0.5, training=self.training)

        x = x.permute(2, 0, 1)
        x, _ = self.mha(x, x, x)
        x = x.permute(1, 2, 0)

        x = self.pool(x).squeeze(2)  # (B, out_ch)

        meta_feat = self.meta_fc(meta)  # (B, out_ch)
        combined = torch.cat([x, meta_feat], dim=1)

        x = self.fc_1(combined)
        out = self.sigmoid(x)

        return out

if __name__ == "__main__":
    # Model parameters
    num_classes = 1      # Number of output classes
    input_channels = 12  # Number of input channels
    seq_length = 2048    # Sequence length
    batch_size = 50     # Batch size
    meta_input = torch.randn(batch_size, 7)  # Meta features

    # Instantiate the ResNet model
    model = MetaResNet1d(nOUT=num_classes, in_ch=input_channels, out_ch=256, mid_ch=64)
    print("ResNet model initialized successfully.\n")

    # Create a sample input tensor (batch_size, channels, seq_length)
    input_tensor = torch.randn(batch_size, input_channels, seq_length)
    print(f"Sample input tensor shape: {input_tensor.shape}\n")

    # Forward pass
    output = model(input_tensor, meta_input)
    print(f"Model output shape: {output.shape}\n")

    # Validate output shape
    expected_shape = (batch_size, num_classes)
    assert output.shape == expected_shape, f"Error: Expected output shape {expected_shape}, but got {output.shape}"
    print("Test passed! The ResNet model works as expected.")